import logging
import numpy as np

from util.util_log import setup_logging
from util.util_csv import save_csv
from util.util_image import save_img, save_img_xyz, draw_img
from util.util_ris_pattern import point_2_phi_pattern, phase_2_pattern, phase_2_pattern_xyz, eps, \
    phase_2_bit, phaseBit_2_deg

from multi_beam.multi_beam_PS import psm_beam_n
from multi_beam.multi_beam_QE import qe_beam_N


# 配置日志，默认打印到控制台，也可以设置打印到文件
setup_logging()
# setup_logging(log_file="../../files/logs/multi_beam_BNPA.log")
# 获取日志记录器并记录日志
logger = logging.getLogger("[RIS-multi-beam-BNPA-1bit]")


# ============================================= 核心方法 ======================================
# 以 center 为中心, angles 为届, 分割 shape
def create_mask(shape, center, angles):
    """创建一个掩码，用于标记由三条直线分割的六个区域"""
    h, w = shape
    mask = np.zeros(shape, dtype=np.uint8)
    # 中心点坐标
    cx, cy = center
    #
    for y in range(h):
        for x in range(w):
            # 计算相对于中心点的相对位置
            dx = x - cx
            dy = y - cy
            # 将直角坐标转换为极坐标角度 (以弧度为单位)
            angle = np.arctan2(dy, dx) * 180 / np.pi
            if angle < 0:
                angle += 360  # 确保角度在 [0, 360] 范围内
            # 确定该像素属于哪个区域
            for i, angle_start in enumerate(angles):
                angle_end = angles[(i + 1) % len(angles)]
                if angle_start <= angle < angle_end:
                    mask[y, x] = i + 1
                    break
            else:
                # 如果没有匹配到任何区域（例如，在中心点上），可以指定为0或忽略
                mask[y, x] = 0
    return mask


# 计算 mask 各值的面积
def calculate_areas(mask):
    """计算每个区域的面积"""
    unique, counts = np.unique(mask, return_counts=True)
    areas = dict(zip(unique, counts))
    # 排除背景（通常是0）
    # if 0 in areas:
    #     del areas[0]
    # 如果有需要的话，可以根据image中的实际内容调整面积计算方式
    return areas


def sort_keys_by_values(dict_area):
    """
    根据字典的值对键进行排序，并返回排序后的键列表。

    参数:
        dict_area (dict): 包含键值对的字典。

    返回:
        list: 按照值从大到小排序后的键列表。
    """
    # 使用 sorted 函数和 lambda 表达式根据值对项进行排序，reverse=True 表示降序排列
    sorted_items = sorted(dict_area.items(), key=lambda item: item[1], reverse=True)
    # 提取排序后的键
    sorted_keys = [item[0] for item in sorted_items]
    return sorted_keys


# 根据array_dict，替换array_mask元素
def mask_replace_multiple(array_mask, array_dict):
    """
    根据array_mask中元素的值来替换元素。

    参数:
        array_mask (np.array): 包含1到N整数的二维numpy数组，指示选择哪个数组。
        array_dict (dict): 包含键1到N分别对应array_1到array_N的字典。

    返回:
        np.array: 替换后的结果数组。
    """
    # 确保输入正确
    if not isinstance(array_mask, np.ndarray) or array_mask.ndim != 2:
        raise ValueError("array_mask 必须是二维的 numpy 数组")
    # 检查 array_dict 是否包含所有需要的键，并且对应的数组形状是否一致
    unique_values = np.unique(array_mask)
    if not all(val in array_dict for val in unique_values):
        raise ValueError(f"array_dict 必须包含键 {unique_values}")
    if not all(isinstance(arr, np.ndarray) and arr.shape == array_mask.shape
               for arr in array_dict.values()):
        raise ValueError("array_dict 中的数组必须是与 array_mask 形状相同的 numpy 数组")
    # 初始化结果数组，可以先用其中一个数组初始化
    result = np.zeros_like(array_mask, dtype=array_dict[1].dtype)
    # 使用布尔索引来替换元素
    for key, array in array_dict.items():
        result[array_mask == key] = array[array_mask == key]
    return result


# 生成填空 mask 的列表
def generate_mask_replace(beam_num, list_area_sort_key, phaseBits, phaseBitDefault):
    dict_replace = {}
    for i in range(len(list_area_sort_key)):
        if i < beam_num:
            dict_replace[list_area_sort_key[i]] = phaseBits[i]
        else:
            dict_replace[list_area_sort_key[i]] = phaseBitDefault
    return dict_replace


# ============================================= 主函数 ====================================
# 核心方法: 双波束
def bnpa_beam_2(phase1, phaseBit1, phase2, phaseBit2, angles, bit_num):
    # 计算混合区域方法
    # phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_beam_n([phase1, phase2], bit_num)    # 使用 PS
    phaseBit_mix, random_indices = qe_beam_N([phaseBit1, phaseBit2])                    # 使用 QE
    # 生成切分码阵的 mask
    rows, cols = phaseBit1.shape
    center = (rows // 2, cols // 2)
    phase_mask = create_mask(phaseBit1.shape, center, angles)
    # draw_img(phase_mask)
    # 计算 mask 的面积并排序
    dict_area = calculate_areas(phase_mask)
    list_area_sort_key = sort_keys_by_values(dict_area)
    # 对应 mask 各值和 phaseBit 的关系 dict_replace
    dict_replace_bit = generate_mask_replace(2, list_area_sort_key, [phaseBit1, phaseBit2], phaseBit_mix)
    # 根据 dict_replace 将 phaseBits 对应位置元素替代分好区的 mask
    phaseBit_BNPA = mask_replace_multiple(phase_mask, dict_replace_bit)
    phaseBitDeg_BNPA = phaseBit_2_deg(phaseBit_BNPA, bit_num)
    # draw_img(phaseBit_BNPA)
    # draw_img(phaseBitDeg_BNPA)
    return phaseBit_BNPA, phaseBitDeg_BNPA


# 核心方法: 双波束
def bnpa_beam_4(phase1, phaseBit1, phase2, phaseBit2,
                phase3, phaseBit3, phase4, phaseBit4, angles, bit_num):
    # 计算混合区域方法
    # phase_mix, phaseBit_mix, phaseBitDeg_mix = psm_beam_n([phase1, phase2, phase3, phase4], bit_num)    # 使用 PS
    phaseBit_mix, random_indices = qe_beam_N([phaseBit1, phaseBit2, phaseBit3, phaseBit4])              # 使用 QE
    # 生成切分码阵的 mask
    rows, cols = phaseBit1.shape
    center = (rows // 2, cols // 2)
    phase_mask = create_mask(phaseBit1.shape, center, angles)
    # draw_img(phase_mask)
    # 计算 mask 的面积并排序
    dict_area = calculate_areas(phase_mask)
    list_area_sort_key = sort_keys_by_values(dict_area)
    # 对应 mask 各值和 phaseBit 的关系 dict_replace
    dict_replace_bit = generate_mask_replace(4, list_area_sort_key,
                                             [phaseBit1, phaseBit2, phaseBit3, phaseBit4], phaseBit_mix)
    # 根据 dict_replace 将 phaseBits 对应位置元素替代分好区的 mask
    phaseBit_BNPA = mask_replace_multiple(phase_mask, dict_replace_bit)
    phaseBitDeg_BNPA = phaseBit_2_deg(phaseBit_BNPA, bit_num)
    # draw_img(phaseBit_BNPA)
    # draw_img(phaseBitDeg_BNPA)
    return phaseBit_BNPA, phaseBitDeg_BNPA


# 相位合成法 -- 双波束
def main_multi_beam_2(theta1, phi1, theta2, phi2, angles, path_pre, bit_num):
    logger.info("main_multi_beam_2: bit_num=%d, path_pre=%s, " % (bit_num, path_pre))
    logger.info("main_multi_beam_2: angles=%s, " % angles)
    logger.info("main_multi_beam_2: theta1=%d, phi1=%d, theta2=%d, phi2=%d, " % (theta1, phi1, theta2, phi2))
    # 目前只支持2bit
    if bit_num > 2:
        logger.error("main_multi_beam_2: bit_num bigger than 2.")
        return
    phase1, phaseBit1, pattern1 = point_2_phi_pattern(theta1, phi1, bit_num)
    phase2, phaseBit2, pattern2 = point_2_phi_pattern(theta2, phi2, bit_num)
    # BNPA 核心方法
    phaseBit_mix, phaseBitDeg_mix = bnpa_beam_2(phase1, phaseBit1, phase2, phaseBit2, angles, bit_num)
    # 计算phase_mix的方向图
    phaseBit_mix = np.deg2rad(phaseBitDeg_mix)
    patternBit_mix = phase_2_pattern(phaseBit_mix)
    patternBit_mix_xyz, x, y, z  = phase_2_pattern_xyz(phaseBit_mix)
    #
    # 保存结果
    logger.info("save BNPA multi-beam 2 result...")
    # 保存图片
    save_img(path_pre + "phase1.jpg", phase1)
    save_img(path_pre + "phase2.jpg", phase2)
    save_img(path_pre + "phaseBit1.jpg", phaseBit1)
    save_img(path_pre + "phaseBit2.jpg", phaseBit2)
    save_img(path_pre + "pattern1.jpg", pattern1)
    save_img(path_pre + "pattern2.jpg", pattern2)
    save_img(path_pre + "phaseBit_mix.jpg", phaseBit_mix)         # BNPA -- 结果码阵
    save_img(path_pre + "patternBit_mix.jpg", patternBit_mix)     # BNPA -- 结果码阵方向图
    save_img_xyz(path_pre + "patternBit_mix_xyz.jpg", np.abs(patternBit_mix_xyz), x, y)
    # 保存相位结果
    save_csv(phase1, path_pre + "phase1.csv")
    save_csv(phase2, path_pre + "phase2.csv")
    save_csv(phaseBit1, path_pre + "phaseBit1.csv")
    save_csv(phaseBit2, path_pre + "phaseBit2.csv")
    save_csv(phaseBit_mix, path_pre + "phaseBit_mix.csv")


# 相位合成法 -- 四波束
def main_multi_beam_4(theta1, phi1, theta2, phi2, theta3, phi3, theta4, phi4, angles, path_pre, bit_num):
    logger.info("main_multi_beam_4: bit_num=%d, path_pre=%s, " % (bit_num, path_pre))
    logger.info("main_multi_beam_4: angles=%s, " % angles)
    logger.info("main_multi_beam_4: theta1=%d, phi1=%d, theta2=%d, phi2=%d, theta3=%d, phi3=%d, theta4=%d, phi4=%d"
                % (theta1, phi1, theta2, phi2, theta3, phi3, theta4, phi4))
    phase1, phaseBit1, pattern1 = point_2_phi_pattern(theta1, phi1, bit_num)
    phase2, phaseBit2, pattern2 = point_2_phi_pattern(theta2, phi2, bit_num)
    phase3, phaseBit3, pattern3 = point_2_phi_pattern(theta3, phi3, bit_num)
    phase4, phaseBit4, pattern4 = point_2_phi_pattern(theta4, phi4, bit_num)
    # 确保所有数组具有相同的形状
    assert phaseBit1.shape == phaseBit2.shape == phaseBit3.shape == phaseBit4.shape, "所有数组必须具有相同的形状"
    # PS
    phaseBit_mix, phaseBitDeg_mix = bnpa_beam_4(phase1, phaseBit1, phase2, phaseBit2, phase3, phaseBit3,
                                                phase4, phaseBit4, angles, bit_num)
    # 计算phase_mix的方向图
    phaseBit_mix = np.deg2rad(phaseBitDeg_mix)
    patternBit_mix = phase_2_pattern(phaseBit_mix)
    #
    # 保存结果
    logger.info("save PS multi-beam 4 result...")
    patternBit_mix_xyz, x, y, z = phase_2_pattern_xyz(phaseBit_mix)
    # 保存图片
    save_img(path_pre + "phase1.jpg", phase1)
    save_img(path_pre + "phase2.jpg", phase2)
    save_img(path_pre + "phase3.jpg", phase3)
    save_img(path_pre + "phase4.jpg", phase4)
    save_img(path_pre + "phaseBit1.jpg", phaseBit1)
    save_img(path_pre + "phaseBit2.jpg", phaseBit2)
    save_img(path_pre + "phaseBit3.jpg", phaseBit3)
    save_img(path_pre + "phaseBit4.jpg", phaseBit4)
    save_img(path_pre + "pattern1.jpg", pattern1)
    save_img(path_pre + "pattern2.jpg", pattern2)
    save_img(path_pre + "pattern3.jpg", pattern3)
    save_img(path_pre + "pattern4.jpg", pattern4)
    save_img(path_pre + "phaseBit_mix.jpg", phaseBit_mix)  # 相位合成法 -- 结果码阵
    save_img(path_pre + "patternBit_mix.jpg", patternBit_mix)  # 相位合成法 -- 结果码阵方向图
    save_img_xyz(path_pre + "patternBit_mix_xyz.jpg", np.abs(patternBit_mix_xyz), x, y)
    # 保存相位结果
    save_csv(phase1, path_pre + "phase1.csv")
    save_csv(phase2, path_pre + "phase2.csv")
    save_csv(phase3, path_pre + "phase3.csv")
    save_csv(phase4, path_pre + "phase4.csv")
    save_csv(phaseBit1, path_pre + "phaseBit1.csv")
    save_csv(phaseBit2, path_pre + "phaseBit2.csv")
    save_csv(phaseBit3, path_pre + "phaseBit3.csv")
    save_csv(phaseBit4, path_pre + "phaseBit4.csv")
    save_csv(phaseBit_mix, path_pre + "phaseBit_mix.csv")



def main_multi_bnpa():
    main_multi_beam_2(30, 0, 30, 90, [0, 90, 180, 270],
                      "../files/multi-beam/2bit/BNPA(QE)/2-(30,0,30,90)/", 2)
    main_multi_beam_4(30, 0, 30, 60, 30, 120, 30, 180, [30, 90, 150, 210, 270, 330],
                      "../files/multi-beam/2bit/BNPA(QE)/4-(30,0,30,60,30,120,30,180)/", 2)



if __name__ == '__main__':
    logger.info("1bit-RIS-multi-beam-BNPA: Beam Normal Partitioning Algorithm")
    main_multi_bnpa()